package rabbit.flt.plugins.reactor.http.plugin.impl;

import io.netty.handler.codec.http.HttpResponseStatus;
import org.reactivestreams.Publisher;
import rabbit.flt.common.trace.TraceData;
import rabbit.flt.common.trace.io.HttpResponse;
import rabbit.flt.plugins.reactor.http.plugin.proxy.ConnectionProxy;
import reactor.core.publisher.Flux;
import reactor.netty.Connection;
import reactor.netty.http.client.HttpClientResponse;

import java.lang.reflect.Method;
import java.util.function.BiFunction;

import static rabbit.flt.plugins.reactor.http.plugin.HttpClientFinalizerPlugin.CONTEXT_KEY_NAME;

public class ResponseConnectionPlugin extends ResponsePlugin {

    @Override
    public Object[] before(Object objectEnhanced, Method method, Object[] args) {
        if (isTraceOpened()) {
            args[0] = proxyResponseConnectionHandler((BiFunction<HttpClientResponse, Connection, Publisher>) args[0]);
        }
        return args;
    }

    /**
     * 代理处理函数
     *
     * @param receiver
     * @return
     */
    private BiFunction<HttpClientResponse, Connection, Publisher> proxyResponseConnectionHandler(BiFunction<HttpClientResponse, Connection, Publisher> receiver) {
        return (httpClientResponse, connection) -> Flux.deferContextual(ctx -> {
            TraceData traceData = ctx.get(CONTEXT_KEY_NAME);
            HttpResponse response = traceData.getHttpResponse();
            httpClientResponse.responseHeaders().forEach(entry -> response.addHeader(entry.getKey(), truncate(entry.getValue())));
            response.setStatusCode(httpClientResponse.status().code());
            traceData.updateCost();
            if (HttpResponseStatus.OK.code() != response.getStatusCode()) {
                Connection proxy = ConnectionProxy.proxy(connection, traceData);
                return Flux.from(receiver.apply(httpClientResponse, proxy));
            }
            return receiver.apply(httpClientResponse, connection);
        });
    }


}
